# -*- coding: utf-8 -*-
# @Time     : 2023/12/14 09:33
# @Author   : Long-Long Qiu
# @FileName : PlotPainter.py
# @Product  : PyCharm
# import packages


import plotly.graph_objects as go
import numpy as np


class PlotPainter(object):
    """
    使用plot库绘制三维图像
    """
    @staticmethod
    def plotCube(fig, x_bottom, y_bottom, z_bottom, length, width, height, color, mode, lineWidth=5, text='', annotations=[]):
        """

        :param fig:
        :param x_bottom: 顶点坐标
        :param y_bottom:
        :param z_bottom:
        :param length:
        :param width:
        :param height:
        :param color:
        :param mode: 1：全透明，2：带颜色
        :param lineWidth:
        :return:
        """
        # 定义立方体的八个顶点
        vertices = [
            [x_bottom, y_bottom, z_bottom],
            [x_bottom + length, y_bottom, z_bottom],
            [x_bottom + length, y_bottom + width, z_bottom],
            [x_bottom, y_bottom + width, z_bottom],
            [x_bottom, y_bottom, z_bottom + height],
            [x_bottom + length, y_bottom, z_bottom + height],
            [x_bottom + length, y_bottom + width, z_bottom + height],
            [x_bottom, y_bottom + width, z_bottom + height]
        ]

        x = [item[0] for item in vertices]
        y = [item[1] for item in vertices]
        z = [item[2] for item in vertices]

        # 绘制每个面
        opacity = 0.6 if mode == 2 else 0.
        fig.add_trace(go.Mesh3d(
            x=x,
            y=y,
            z=z,
            i=[0, 0, 4, 4, 1, 5, 0, 4, 1, 1, 2, 6],
            j=[1, 2, 5, 6, 2, 2, 3, 3, 5, 4, 3, 3],
            k=[2, 3, 6, 7, 5, 6, 4, 7, 4, 0, 6, 7],
            color=color,
            opacity=opacity,
            lighting=dict(ambient=0.8),
            showlegend=False
        ))

        # 绘制立方体的边框
        x_edges = [x[0] for x in vertices]
        y_edges = [x[1] for x in vertices]
        z_edges = [x[2] for x in vertices]

        # 绘制边框线
        def add_edge(i, j):
            fig.add_trace(go.Scatter3d(
                x=[x_edges[i], x_edges[j]],
                y=[y_edges[i], y_edges[j]],
                z=[z_edges[i], z_edges[j]],
                mode='lines',
                line=dict(color=color, width=4),
                showlegend=False
            ))

        # 绘制边框
        add_edge(0, 1)
        add_edge(1, 2)
        add_edge(2, 3)
        add_edge(3, 0)
        add_edge(4, 5)
        add_edge(5, 6)
        add_edge(6, 7)
        add_edge(7, 4)
        add_edge(0, 4)
        add_edge(1, 5)
        add_edge(2, 6)
        add_edge(3, 7)

        # 绘制文字
        if text != '':
            # 添加文字注释到注释列表
            annotations.append(dict(
                showarrow=True,
                x=x_bottom + length / 2,
                y=y_bottom + width / 2,
                z=z_bottom + height / 2,
                text=text,
                font=dict(color='black')
            ))

            # 将列表转换回元组并更新图表的注释列表
            fig.update_layout(scene=dict(annotations=tuple(annotations)))

        return fig

    @staticmethod
    def slice_triangles(z, n, i, j, k, l):
        """Create the triangles of a single slice"""
        return [[z, j, i], [i, j, l], [l, j, k], [k, n, l]]

    @staticmethod
    def cylinder_mesh(r, xs, ys, zs, h, n_slices=40):
        """Create a cylindrical mesh"""
        theta = np.linspace(0, 2 * np.pi, n_slices + 1)
        x = xs + r * np.cos(theta)
        z = zs + r * np.sin(theta)
        y1 = ys + 0 * np.ones_like(x)
        y2 = (ys + h) * np.ones_like(x)

        # index of the final point in the mesh
        n = n_slices * 2 + 1

        # build triangulation
        triangles = []
        for s in range(1, n_slices + 1):
            j = (s + 1) if (s <= n_slices - 1) else 1
            k = j + n_slices if (s <= n_slices - 1) else n_slices + 1
            l = s + n_slices
            triangles += PlotPainter.slice_triangles(0, n, s, j, k, l)
        triangles = np.array(triangles)

        # coordinates of the vertices
        x_coords = np.hstack([xs, x[:-1], x[:-1], xs])
        z_coords = np.hstack([zs, z[:-1], z[:-1], zs])
        y_coords = np.hstack([ys, y1[:-1], y2[:-1], (ys + h)])
        vertices = np.stack([x_coords, y_coords, z_coords]).T

        return vertices, triangles, x, z, y1, y2

    @staticmethod
    def circle_parametric_equation(radius, theta, y_coord, x_coord, z_coord):
        cx = radius * np.cos(theta) + x_coord
        cz = radius * np.sin(theta) + z_coord
        cy = np.ones_like(theta) * y_coord
        return cx, cy, cz

    @staticmethod
    def plotCylinder(fig, r, xs, ys, zs, h, n_slices=40, show_mesh=True, n_sub=4, surface_kw={}, line_kw={}, paint_grid=False):
        """
            r : radius
            xs, ys, zs : start position of the cylinder
            h : height of the cylinder
            n_slices : number of slices in the circumferential direction
            show_mesh : wheter to display pseudo-wireframe
            n_sub : number of subdivision in along the height for the pseudo-wireframe
            surface_kw : customize the appearance of the surface
            line_kw : customize the appearance of the wireframe
            """
        vertices, triangles, x, z, y1, y2 = PlotPainter.cylinder_mesh(r, xs, ys, zs, h, n_slices)
        surface = go.Mesh3d(
            x=vertices[:, 0], y=vertices[:, 1], z=vertices[:, 2],
            i=triangles[:, 0], j=triangles[:, 1], k=triangles[:, 2],
            **surface_kw)

        traces = [surface]
        if not show_mesh:
            return traces

        line_kw.setdefault("showlegend", False)

        if paint_grid:
            # horizontal mesh lines
            zsubs = np.linspace(zs, zs + h, n_sub + 1)
            padding = 0
            for zc in zsubs:
                traces.append(go.Scatter3d(x=x, z=z, y=zc * np.ones_like(x) + padding, mode="lines", **line_kw))
            # vertical mesh lines
            for _x, _y in zip(x, z):
                traces.append(go.Scatter3d(x=[_x, _x], z=[_y, _y], y=[zs + padding, zs + h + padding], mode="lines", **line_kw))

        # 绘制圆柱体的两个底圆
        theta = np.linspace(0, 2 * np.pi, 300)

        # 第一个底圆
        # 计算圆的坐标
        x1, y1, z1 = PlotPainter.circle_parametric_equation(r, theta, ys, xs, zs)
        traces.append(go.Scatter3d(x=x1, y=y1, z=z1, mode='lines', **line_kw))

        # 第二个底圆
        # 计算圆的坐标
        x2, y2, z2 = PlotPainter.circle_parametric_equation(r, theta, ys + h, xs, zs)
        traces.append(go.Scatter3d(x=x2, y=y2, z=z2, mode='lines', **line_kw))

        fig.add_traces(traces)
        return fig


if __name__ == '__main__':
    # 创建图形对象
    fig = go.Figure()
    PlotPainter.plotCube(fig, 0, 0, 0, 3, 4, 5, 'green', 2)
    PlotPainter.plotCylinder(fig, 2, 4, 4, 2, 4, 20, n_sub=4, line_kw={"line_color": "#202020", "line_width": 3}, surface_kw={"color": "red"})
    # 设置图表布局
    fig.update_layout(scene=dict(
        xaxis=dict(visible=True),
        yaxis=dict(visible=True),
        zaxis=dict(visible=True),
        aspectmode='cube'  # 设置刻度相等
    ))
    fig.show()
